import random
import math

import cv2


class Node:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.parent = None

# 生成随机点
def get_random_node(width, height):
    return Node(random.randint(0, width), random.randint(0, height))


# 求距离
def distance(node1, node2):
    return math.sqrt((node1.x - node2.x) ** 2 + (node1.y - node2.y) ** 2)

# 找到距离最近的节点
def get_nearest_node(node_list, random_node):
    return min(node_list, key=lambda node: distance(node, random_node))


def rrt_path_detect(image_path, start, goal, passwall=False, max_iter=1000):
    # 加载图片
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    height, width = image.shape
    # 生成图
    color_image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

    # 开始节点
    start_node = Node(start[0], start[1])
    # 目标节点
    goal_node = Node(goal[0], goal[1])

    # node列表
    node_list = [start_node]
    cv2.circle(color_image, (start_node.x, start_node.y), 1, (0, 0, 0), -1)  # 半径为1，颜色为绿色，-1表示填充

    #迭代总次数
    num = 0

    for _ in range(max_iter):
        # 统计总次数
        num = num+1
        # 生成随机点
        random_node = get_random_node(width, height)
        # 获取最近的节点
        nearest_node = get_nearest_node(node_list, random_node)

        # 生成new节点
        theta = math.atan2(random_node.y - nearest_node.y, random_node.x - nearest_node.x)
        new_node = Node(
            int(nearest_node.x + 10 * math.cos(theta)),
            int(nearest_node.y + 10 * math.sin(theta))
        )
        # 防止溢出
        new_node.x = min(new_node.x, width-1)
        new_node.x = max(new_node.x, 0)
        new_node.y = min(new_node.y, height-1)
        new_node.y = max(new_node.y, 0)
        new_node.parent = nearest_node
        # 画点
        cv2.circle(color_image, (new_node.x, new_node.y), 1, (0, 0, 0), -1)  # 半径为1，颜色为绿色，-1表示填充
        # 画线
        cv2.line(color_image, (nearest_node.x, nearest_node.y), (new_node.x, new_node.y),
                 (0, 255, 0), 1)

        if passwall:
            node_list.append(new_node)
        elif image[new_node.y, new_node.x] >= 249:  # 不能是障碍物
            node_list.append(new_node)

        if distance(new_node, goal_node) < 10:
            goal_node.parent = new_node
            node_list.append(goal_node)
            print("找到了，通过"+str(num)+"次")
            break

    # 画路线
    current_node = goal_node
    while current_node.parent is not None:
        cv2.line(color_image, (current_node.x, current_node.y), (current_node.parent.x, current_node.parent.y),
                 (0, 0, 0), 2)
        current_node = current_node.parent

    cv2.imshow('RRT Path', color_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # 使用该函数
    rrt_path_detect('path_to_your_image.jpg', (start_x, start_y), (goal_x, goal_y))
